﻿module Physics.Chapter4.CollisionResponse

open Casanova
open Casanova.Input
open Casanova.Core
open Casanova.Coroutines
open Casanova.Utilities
open Casanova.Math
open Casanova.Game
open Casanova.Drawing
open Casanova.StandardLibrary
open Casanova.StandardLibrary.Core
open Physics.Chapter4.RigidBody
open Physics.Chapter4.Algebra

type CollisionConstraint = {
  Body1               : Ref<RigidBody>
  Body2               : Ref<RigidBody>
  N                   : Vector3<1>
  P                   : Vector3<m>
  Depth               : float32<m>
}

let mutable private prev_x = [||]

let collision_response (collision_constraints:List<CollisionConstraint>)
                       (bodies:RigidBody[])
                       (dt:float32<s>) =
  if dt < 1.0f<s> / 30.0f then
    let dt_left = 900.0f * (1.0f<s> / 30.0f - dt) |> int
    System.Threading.Thread.Sleep(dt_left)
  let dt = 1.0f<s> / 30.0f

  // tuples of the type i,j (i < j), a_i, b_i, a_j, b_j, k
  // a_i, b_i are the constraint vectors for body i
  // a_j, b_j are the constraint vectors for body j
  // k is the pushing factor, to allow the force to 
  //   perform work and restore the constraint if it 
  //   is broken
  let constraints = 
    [|
      for c in collision_constraints do
        let b_i,b_j = !c.Body1, !c.Body2

        if !b_i.Asleep && b_i.Immovable |> not then
          do b_i.Asleep := false
        if !b_j.Asleep && b_j.Immovable |> not then
          do b_j.Asleep := false

        let x_i,x_j = !b_i.Position, !b_j.Position
        let r_i,r_j = c.P - x_i, c.P - x_j
        let n = c.N * 1.0f<_>
        let k = 
          if abs c.Depth < 0.01f<m> then 0.001f 
          else c.Depth * 0.1f<_>

        let j1,j2 = -n, -Vector3<1>.Cross(r_i, n)
        let j3,j4 =  n,  Vector3<1>.Cross(r_j, n)

        let j1,j2,j3,j4 = 
          if b_i.MassInv = 0.0f<_> then
            Vector3.Zero, Vector3.Zero, j3, j4
          elif b_j.MassInv = 0.0f<_> then
            j1, j2, Vector3.Zero, Vector3.Zero
          else
            j1, j2, j3, j4

        yield b_i, b_j, 
              j1, j2,
              j3, j4,
              k,
              -1000.0f, 0.0f
    |]

  for i = 0 to bodies.Length - 1 do
    let b = bodies.[i]
    do b.Index <- i
    if (!b.F).Length > 0.0f<N> || (!b.T).Length > 0.0f<N * m> then
      b.Asleep := false
    if !b.Asleep |> not && b.MassInv <> 0.0f<kg^-1> then
      let G = Vector3<N>.Create(0.0f, float32(-9.81f / b.MassInv), 0.0f)
      b.F := !b.F + G

  if constraints.Length = 0 then 
    for b_i in bodies do
      if b_i.Immovable |> not then
        let v2 = !b_i.Velocity + b_i.MassInv * !b_i.F * dt
        let w2 = !b_i.AngularVelocity + b_i.InverseI * !b_i.T * dt * 1.0f<_>
        b_i.Velocity := v2
        b_i.AngularVelocity := w2
        if (!b_i.Velocity).Length <= 0.001f<m/s> && (!b_i.AngularVelocity).Length <= 0.001f<m^2/s> then
          b_i.Asleep := true
        b_i.Position := !b_i.Position + !b_i.Velocity * dt
        b_i.Rotation := !b_i.Rotation + b_i.RotationVelocity * (float32 dt)
        b_i.Rotation := Microsoft.Xna.Framework.Quaternion.Normalize(!b_i.Rotation)
  else
    let b = 
      [|
        for b_i,b_j,dv_i,dw_i,dv_j,dw_j,k,x_min,x_max in constraints do
          let e = float32 k
          let v_i,w_i,v_j,w_j = !b_i.Velocity, !b_i.AngularVelocity, !b_j.Velocity, !b_j.AngularVelocity
          let jv = Vector3<1>.Dot(v_i, dv_i) + Vector3<1>.Dot(w_i, dw_i * 1.0f<_>) +
                    Vector3<1>.Dot(v_j, dv_j) + Vector3<1>.Dot(w_j, dw_j * 1.0f<_>)

          let f_i,t_i,f_j,t_j = !b_i.F, !b_i.T, !b_j.F, !b_j.T
          let a_i,al_i,a_j,al_j = f_i * b_i.MassInv, b_i.InverseI * t_i,
                                  f_j * b_j.MassInv, b_j.InverseI * t_j
          let jmf = Vector3<1>.Dot(dv_i, a_i) + Vector3<1>.Dot(dw_i, al_i * 1.0f<_>) +
                      Vector3<1>.Dot(dv_j, a_j) + Vector3<1>.Dot(dw_j, al_j * 1.0f<_>)
          yield e / dt * 1.0f<_> - jv * 1.0f<_> - 1.0f<_> * jmf * dt
      |] |> Vector.Create

    let JMInverse = 
      let n = bodies.Length
      [|
        for b_i,b_j,dv_i,dw_i,dv_j,dw_j,k,x_min,x_max in constraints do
          let v_i,w_i,v_j,w_j = !b_i.Velocity, !b_i.AngularVelocity, !b_j.Velocity, !b_j.AngularVelocity

          let f_i,t_i,f_j,t_j = !b_i.F, !b_i.T, !b_j.F, !b_j.T
          let a_i,al_i,a_j,al_j = f_i * b_i.MassInv, b_i.InverseI * t_i,
                                  f_j * b_j.MassInv, b_j.InverseI * t_j
          let jmf = Vector3<1>.Dot(dv_i, a_i) + Vector3<1>.Dot(dw_i, al_i * 1.0f<_>) +
                      Vector3<1>.Dot(dv_j, a_j) + Vector3<1>.Dot(dw_j, al_j * 1.0f<_>)

          yield b_i.Index,b_j.Index,dv_i * b_i.MassInv,b_i.InverseI * dw_i,dv_j * b_j.MassInv,b_j.InverseI * dw_j
      |]

    let A = 
      let m = constraints.Length    
      [|
        for x = 0 to m - 1 do
          yield 
            [|
              for y = 0 to m - 1 do

                // sum row x of JMInverse and row y of J
                let i0_x,i1_x,a_i0_x,b_i0_x,a_i1_x,b_i1_x = JMInverse.[x]
                let i0_y,i1_y,a_i0_y,b_i0_y,a_i1_y,b_i1_y,_,_,_ = constraints.[y]
                let i0_y,i1_y = i0_y.Index,i1_y.Index

                let sum = var 0.0f
                let inline check_and_sum (i_x,a_i_x,b_i_x) (i_y,a_i_y,b_i_y) =
                    if i_x = i_y then
                      sum := !sum + Vector3<1>.Dot(a_i_x, a_i_y) +
                                    Vector3<1>.Dot(b_i_x, b_i_y)

                let (i_x,a_i_x,b_i_x) = (i0_x,a_i0_x,b_i0_x)
                let (i_y,a_i_y,b_i_y) = (i0_y,a_i0_y,b_i0_y)
                do check_and_sum (i_x,a_i_x,b_i_x) (i_y,a_i_y * 1.0f<_>,b_i_y * 1.0f<_>)
                let (i_y,a_i_y,b_i_y) = (i1_y,a_i1_y,b_i1_y)
                do check_and_sum (i_x,a_i_x,b_i_x) (i_y,a_i_y * 1.0f<_>,b_i_y * 1.0f<_>)

                let (i_x,a_i_x,b_i_x) = (i1_x,a_i1_x,b_i1_x)
                let (i_y,a_i_y,b_i_y) = (i0_y,a_i0_y,b_i0_y)
                do check_and_sum (i_x,a_i_x,b_i_x) (i_y,a_i_y * 1.0f<_>,b_i_y * 1.0f<_>)
                let (i_y,a_i_y,b_i_y) = (i1_y,a_i1_y,b_i1_y)
                do check_and_sum (i_x,a_i_x,b_i_x) (i_y,a_i_y * 1.0f<_>,b_i_y * 1.0f<_>)
                  
                yield !sum
            |]
      |]
    let A = Matrix.Create(A)

    let x0 = 
      if prev_x.Length <> constraints.Length then
        [| for c in constraints do yield 0.0f |] |> Algebra.Vector.Create
      else
        prev_x |> Algebra.Vector.Create
    let x_min = [| for i,j,a_i,b_i,a_j,b_j,k,x_min,x_max in constraints do yield x_min |] |> Algebra.Vector.Create
    let x_max = [| for i,j,a_i,b_i,a_j,b_j,k,x_min,x_max in constraints do yield x_max |] |> Algebra.Vector.Create
    let lambda = projected_gauss_seidel A b x0 x_min x_max
    do prev_x <- lambda.V

    // let Fc = J.Transpose * lambda
    let Fc = 
      [|
        for b in bodies do
          yield Vector3.Zero
      |]
    let Tc = 
      [|
        for b in bodies do
          yield Vector3.Zero
      |]

    for (i,j,a_i,b_i,a_j,b_j,_,_,_),lambda in Seq.zip constraints lambda.V do
      let i,j = i.Index,j.Index
      let F_i,T_i = a_i * lambda, b_i * lambda
      let F_j,T_j = a_j * lambda, b_j * lambda
      Fc.[i] <- Fc.[i] + F_i
      Fc.[j] <- Fc.[j] + F_j
      Tc.[i] <- Tc.[i] + T_i
      Tc.[j] <- Tc.[j] + T_j
      
    let V2 = 
      [|
        for b in bodies do
          let Fc,Tc = Fc.[b.Index],Tc.[b.Index]
          let F = !b.F * dt + Fc * 1.0f<_>
          let T = !b.T * dt + Tc * 1.0f<_>
          let v2 = !b.Velocity + b.MassInv * F
          let w2 = !b.AngularVelocity + b.InverseI * T * 1.0f<_>

          yield! v2 |> vec_to_seq
          yield! w2 |> vec_to_seq
      |] |> Vector.Create

    for i = 0 to bodies.Length - 1 do
      let b_i = bodies.[i]
      let v = Vector3(V2.[i * 3 * 2 + 0],
                      V2.[i * 3 * 2 + 1],
                      V2.[i * 3 * 2 + 2])
      let w = Vector3(V2.[i * 3 * 2 + 3],
                      V2.[i * 3 * 2 + 4],
                      V2.[i * 3 * 2 + 5])
      b_i.Velocity := v * 1.0f<_>
      b_i.AngularVelocity := w * 1.0f<_>
      if (!b_i.Velocity).Length <= 0.1f<m/s> && (!b_i.AngularVelocity).Length <= 0.1f<m^2/s> then
        b_i.Asleep := true
      else
        if b_i.MassInv <> 0.0f<_> then
          b_i.Position := !b_i.Position + !b_i.Velocity * dt
          b_i.Rotation := !b_i.Rotation + b_i.RotationVelocity * (float32 dt)
          b_i.Rotation := Microsoft.Xna.Framework.Quaternion.Normalize(!b_i.Rotation)
